import json
import os
import sys

import fire
import torch
import torch.distributed
from zkl_aiutils_datasets import load_dataset
from zkl_serialization import dump_json_value, load_and_parse_json

project_dir_path = os.path.join(os.path.dirname(__file__), "../..")
sys.path.append(project_dir_path)

from llmpt.model import GPTLaunchingHparams, GPTTrainingCreateArgs, GPTTrainingSimple, make_training_name
from scripts.config import default_dataset_path, default_hparams_file_path, trainings_dir_path


def main(*,
    trainings_dir_path: str = trainings_dir_path,
    training_dir_path: str | None = None,
    hparams_file_path: str = default_hparams_file_path,
    dataset_path: str = default_dataset_path,
):
    # training dir
    hparams = load_and_parse_json(hparams_file_path, GPTLaunchingHparams)
    if training_dir_path is None:
        training_dir_path = os.path.join(trainings_dir_path, make_training_name(hparams))
        training_dir_path = os.path.abspath(training_dir_path)

    # print useful information
    print("Starting training:", file=sys.stderr)
    print("training_dir_path=" + training_dir_path, file=sys.stderr)
    print("model_hparams=" + json.dumps(dump_json_value(hparams.model), indent=4), file=sys.stderr)
    print("training_hparams=" + json.dumps(dump_json_value(hparams.training), indent=4), file=sys.stderr)
    print(file=sys.stderr)

    # device configurations
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    compile = True if torch.cuda.is_available() else False
    torch.set_float32_matmul_precision('medium')

    # perform training
    GPTTrainingSimple(GPTTrainingCreateArgs(
        training_dir_path=training_dir_path,
        model_hparams=hparams.model,
        training_hparams=hparams.training,
        train_dataset_factory=lambda: load_dataset(os.path.join(dataset_path, 'train')),
        valid_dataset_factory=lambda: load_dataset(os.path.join(dataset_path, 'valid')),
        device=device,
        compile=compile,
    )).run()


if __name__ == '__main__':
    fire.Fire(main)
